# Copyright (c) 2024, NVIDIA CORPORATION.

from typing import Any

from typing_extensions import Self

from rmm.pylibrmm.stream import Stream

from pylibcudf.io.types import (
    CompressionType,
    SinkInfo,
    SourceInfo,
    StatisticsFreq,
    TableInputMetadata,
    TableWithMetadata,
)
from pylibcudf.table import Table

class OrcReaderOptions:
    def set_num_rows(self, nrows: int) -> None: ...
    def set_skip_rows(self, skip_rows: int) -> None: ...
    def set_stripes(self, stripes: list[list[int]]) -> None: ...
    def set_decimal128_columns(self, val: list[str]) -> None: ...
    def set_columns(self, col_names: list[str]) -> None: ...
    @staticmethod
    def builder(source: SourceInfo) -> OrcReaderOptionsBuilder: ...

class OrcReaderOptionsBuilder:
    def use_index(self, use: bool) -> Self: ...
    def build(self) -> OrcReaderOptions: ...

def read_orc(
    options: OrcReaderOptions, stream: Stream = None
) -> TableWithMetadata: ...

class OrcColumnStatistics:
    def __init__(self): ...
    @property
    def number_of_values(self) -> int | None: ...
    @property
    def has_null(self) -> bool | None: ...
    def __getitem__(self, item: str) -> Any: ...
    def __contains__(self, item: str) -> bool: ...
    def get[T](self, item: str, default: None | T = None) -> T | None: ...

class ParsedOrcStatistics:
    def __init__(self): ...
    @property
    def column_names(self) -> list[str]: ...
    @property
    def file_stats(self) -> list[OrcColumnStatistics]: ...
    @property
    def stripes_stats(self) -> list[OrcColumnStatistics]: ...

def read_parsed_orc_statistics(
    source_info: SourceInfo,
) -> ParsedOrcStatistics: ...

class OrcWriterOptions:
    def __init__(self): ...
    def set_stripe_size_bytes(self, size_bytes: int) -> None: ...
    def set_stripe_size_rows(self, size_rows: int) -> None: ...
    def set_row_index_stride(self, stride: int) -> None: ...
    @staticmethod
    def builder(sink: SinkInfo, table: Table) -> OrcWriterOptionsBuilder: ...

class OrcWriterOptionsBuilder:
    def __init__(self): ...
    def compression(self, comp: CompressionType) -> Self: ...
    def enable_statistics(self, val: StatisticsFreq) -> Self: ...
    def key_value_metadata(self, kvm: dict[str, str]) -> Self: ...
    def metadata(self, meta: TableWithMetadata) -> Self: ...
    def build(self) -> OrcWriterOptions: ...

def write_orc(options: OrcWriterOptions, stream: Stream = None) -> None: ...

class OrcChunkedWriter:
    def __init__(self): ...
    def close(self) -> None: ...
    def write(self, table: Table) -> None: ...
    @staticmethod
    def from_options(
        options: ChunkedOrcWriterOptions, stream: Stream = None
    ) -> Self: ...

class ChunkedOrcWriterOptions:
    def __init__(self): ...
    def set_stripe_size_bytes(self, size_bytes: int) -> None: ...
    def set_stripe_size_rows(self, size_rows: int) -> None: ...
    def set_row_index_stride(self, stride: int) -> None: ...
    @staticmethod
    def builder(sink: SinkInfo) -> ChunkedOrcWriterOptionsBuilder: ...

class ChunkedOrcWriterOptionsBuilder:
    def __init__(self): ...
    def compression(self, comp: CompressionType) -> Self: ...
    def enable_statistics(self, val: StatisticsFreq) -> Self: ...
    def key_value_metadata(self, kvm: dict[str, str]) -> Self: ...
    def metadata(self, meta: TableInputMetadata) -> Self: ...
    def build(self) -> ChunkedOrcWriterOptions: ...
